package edu.zju.gis.dbfg.model.statistic;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFunction;
import org.geotools.referencing.CRS;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.io.FileWriter;
import java.util.Map;

import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_CRS;

public class StatisticByValue {
    private final static Logger logger = LoggerFactory.getLogger(StatisticByValue.class);

    public static void main(String[] args) throws Exception {
        // Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
//        conf.setMaster("local[4]");
//        conf.setAppName("bob");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        //Setup model arguments
        logger.info("Setup arguments");
        StatisticByValueArgs mArgs = StatisticByValueArgs.initArgs(args);
        if (mArgs == null) {
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(StatisticByValue.class, "main()", msg, args);
        }

        CoordinateReferenceSystem crs = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs()));
        CoordinateReferenceSystem targetCRS = DEFAULT_CRS();
        MathTransform transform = null;
        if (!crs.equals(targetCRS)) {
            transform = CRS.findMathTransform(crs, targetCRS);
        }

        String input = mArgs.getInput();
        String type = mArgs.getType();
        if(type.equals("WKT") && !type.startsWith("file://")){
            String preLocalFile = "file://";
            input = String.format("%s%s",preLocalFile, input);
            logger.info(input);
        }
        String summarizeIndex = mArgs.getSummarizeIndex();
        String areaIndex = mArgs.getAreaIndex();

        JavaRDD<GeoItem> gi1 = null;
        if (type.equals("ES")) {
            gi1 = DataReaderRDD.readToRDD2(type, jsc, input, null, null, null, transform);
        } else {
            String idIndex = mArgs.getIdIndex();
            gi1 = DataReaderRDD.readToRDD(type, jsc, input, idIndex, null, transform);
        }

        JavaPairRDD<String, GeoItem> temp = gi1.mapToPair(new PairFunction<GeoItem, String, GeoItem>() {
            @Override
            public Tuple2<String, GeoItem> call(GeoItem input) throws Exception {
                String summarize = String.valueOf(input.getProp().get(summarizeIndex));
                return new Tuple2<>(summarize, input);
            }
        });

        JavaPairRDD<String, GeoItem> reduceResult = temp.reduceByKey(new Function2<GeoItem, GeoItem, GeoItem>() {
            @Override
            public GeoItem call(GeoItem geoi1, GeoItem geoi2) throws Exception {
                Double area = Double.valueOf(String.valueOf(geoi1.getProp().get(areaIndex))) + Double.valueOf(String.valueOf(geoi2.getProp().get(areaIndex)));
                GeoItem result = geoi1;
                result.getProp().put(areaIndex, area);
                return result;
            }
        });

        Map<String, GeoItem> result = reduceResult.collectAsMap();
        File file = new File(mArgs.getOutput());
        if (file.exists()) {
            file.delete();
        }
        FileWriter writer = new FileWriter(file);
        writer.write(String.format("%s\t%s\n", summarizeIndex, areaIndex));
        for (String key : result.keySet()) {
            String summarize = key;
            GeoItem gi = result.get(key);
            Double area = Double.valueOf(String.valueOf(gi.getProp().get(areaIndex)));
            String s = String.format("%s\t%.9f", summarize, area);
            writer.write(s + "\n");
        }
        writer.flush();
        writer.close();
    }

}
